import time
from options.test_options import TestOptions
from data.data_loader_test import CreateDataLoader
from models.networks import ResUnetGenerator, load_checkpoint
from models.afwm import AFWM
import torch.nn as nn
import os
import numpy as np
import torch
import cv2
import torch.nn.functional as F


def main():
    opt = TestOptions().parse()

    start_epoch, epoch_iter = 1, 0

    data_loader = CreateDataLoader(opt)
    dataset = data_loader.load_data()
    dataset_size = len(data_loader)
    print("***21", dataset_size)

    warp_model = AFWM(opt, 3)
    print(warp_model)
    warp_model.eval()
    warp_model.cuda()
    load_checkpoint(warp_model, opt.warp_checkpoint)

    gen_model = ResUnetGenerator(7, 4, 5, ngf=64, norm_layer=nn.BatchNorm2d)
    print(gen_model)
    gen_model.eval()
    gen_model.cuda()
    load_checkpoint(gen_model, opt.gen_checkpoint)

    total_steps = (start_epoch - 1) * dataset_size + epoch_iter
    step = 0
    step_per_batch = dataset_size / opt.batchSize
    for epoch in range(1, 2):

        for i, data in enumerate(dataset, start=epoch_iter):
            t1 = time.time()

            iter_start_time = time.time()

            total_steps += opt.batchSize
            epoch_iter += opt.batchSize

            real_image = data['image']
            clothes = data['clothes']
            # edge is extracted from the clothes image with the built-in function in python
            edge = data['edge']
            edge = torch.FloatTensor((edge.detach().numpy() > 0.5).astype(np.int))
            clothes = clothes * edge

            flow_out = warp_model(real_image.cuda(), clothes.cuda())
            warped_cloth, last_flow, = flow_out
            warped_edge = F.grid_sample(edge.cuda(),
                                        last_flow.permute(0, 2, 3, 1),
                                        mode='bilinear',
                                        padding_mode='zeros')
            # print("***** 61", warped_edge.shape)  # [1, 1, 256, 192]
            gen_inputs = torch.cat([real_image.cuda(), warped_cloth, warped_edge], 1)
            gen_outputs = gen_model(gen_inputs)  # torch.Size([1, 4, 256, 192])
            # print("***** 64", gen_outputs.shape)
            # tensor结构会一共切分成len(list)这么多的小块
            p_rendered, m_composite = torch.split(gen_outputs, [3, 1], 1)  # 切成两块
            p_rendered = torch.tanh(p_rendered)  # 呈现[1, 3, 256, 192]
            m_composite = torch.sigmoid(m_composite)  # 合成的，[1, 1, 256, 192]

            m_composite = m_composite * warped_edge  # [1, 1, 256, 192]，[1, 1, 256, 192] 轮廓卷曲
            p_tryon = warped_cloth * m_composite + p_rendered * (1 - m_composite)  # 最后穿上的模样

            path = 'results/' + opt.name
            os.makedirs(path, exist_ok=True)
            sub_path = path + '/PFAFN'
            os.makedirs(sub_path, exist_ok=True)
            print(time.time()-t1)
            # if step % 1 == 0:
            #     a = real_image.float().cuda()  # 模特
            #     b = clothes.cuda()  # 衣服
            #     c = p_tryon  # 穿上后
            #     combine = torch.cat([a[0], b[0], c[0]], 2).squeeze()
            #     cv_img = (combine.permute(1, 2, 0).detach().cpu().numpy() + 1) / 2
            #     rgb = (cv_img * 255).astype(np.uint8)
            #     bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
            #     cv2.imwrite(sub_path + '/' + str(step) + '.jpg', bgr)
            #     cv2.imshow("24", bgr)
            #     cv2.waitKey(2000)
            #
            # step += 1
            # if epoch_iter >= dataset_size:
            #     break


if __name__ == '__main__':
    main()
    pass
